package com.peng.leetcode.tree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * CountPairs
 * <p>
 * 1530. 好叶子节点对的数量
 *
 * @author: lupeng6
 * @create: 2021/2/4 18:28
 */
public class CountPairs {

    public static void main(String[] args) {
        TreeNode root = TreeMain.array2tree(new Integer[]{1, 2, 3, 4, 5, 6, 7});
        System.out.println(new CountPairs().countPairs(root, 3));
    }

    Map<TreeNode, TreeNode> parent = new HashMap<>();

    List<TreeNode> leaves = new ArrayList<>();

    public int countPairs(TreeNode root, int distance) {
        initParentAndLeaves(root);

        int answer = 0;

        int size = leaves.size();
        for (int i = 0; i < size - 1; i++) {
            for (int j = i + 1; j < size; j++) {
                int length = computePathLength(leaves.get(i), leaves.get(j));
                if (length > 0 && length <= distance) {
                    answer++;
                }
            }
        }

        return answer;
    }

    int computePathLength(TreeNode node1, TreeNode node2) {
        TreeNode curr = node1;
        Map<TreeNode, Integer> pathLen = new HashMap<>();
        while (parent.containsKey(curr)) {
            TreeNode parent = this.parent.get(curr);
            Integer len = pathLen.get(curr);

            pathLen.put(parent, len == null ? 1 : len + 1);
            curr = parent;
        }

        int node2PathLen = 0;
        while (parent.containsKey(node2)) {
            node2PathLen++;
            TreeNode parent = this.parent.get(node2);
            Integer path = pathLen.get(parent);
            if (path != null) {
                return path + node2PathLen;
            }
            node2 = parent;
        }
        return -1;
    }

    void initParentAndLeaves(TreeNode node) {
        if (node.left == null && node.right == null) {
            leaves.add(node);
            return;
        }
        if (node.left != null) {
            parent.put(node.left, node);
            initParentAndLeaves(node.left);
        }
        if (node.right != null) {
            parent.put(node.right, node);
            initParentAndLeaves(node.right);
        }
    }
}
